#!/usr/bin/env python
"""

   translate_by_codon.py

        accepts cDNA sequence and codon by codon translation.
        Unknown codons are translated as 'X'


"""


stops = dict.fromkeys(['TGA', 'TAA', 'TAG'])


def create_aa_nucl_map():
    aa_nucl_map = dict()

    def four_fold_deg (codon12, aa):
        for third in 'A', 'C', 'G', 'T', 'X', 'N', '-', ' ':
            aa_nucl_map[codon12 + third] = aa

    # stops
    aa_nucl_map['TGA'] = '*'
    aa_nucl_map['TAA'] = '*'
    aa_nucl_map['TAG'] = '*'

    four_fold_deg('GC', 'A')

    aa_nucl_map['TGC'] = 'C'
    aa_nucl_map['TGT'] = 'C'
    aa_nucl_map['TGT'] = 'C'
    aa_nucl_map['TGY'] = 'C'

    aa_nucl_map['GAC'] = 'D'
    aa_nucl_map['GAT'] = 'D'
    aa_nucl_map['GAY'] = 'D'

    aa_nucl_map['GAA'] = 'E'
    aa_nucl_map['GAG'] = 'E'
    aa_nucl_map['GAR'] = 'E'

    aa_nucl_map['TTC'] = 'F'
    aa_nucl_map['TTT'] = 'F'
    aa_nucl_map['TTY'] = 'F'

    four_fold_deg('GG', 'G')

    aa_nucl_map['CAC'] = 'H'
    aa_nucl_map['CAT'] = 'H'
    aa_nucl_map['CAY'] = 'H'

    aa_nucl_map['ATA'] = 'I'
    aa_nucl_map['ATC'] = 'I'
    aa_nucl_map['ATT'] = 'I'
    aa_nucl_map['ATH'] = 'I'

    aa_nucl_map['AAA'] = 'K'
    aa_nucl_map['AAG'] = 'K'
    aa_nucl_map['AAG'] = 'K'

    four_fold_deg('CT', 'L')
    aa_nucl_map['TTA'] = 'L'
    aa_nucl_map['TTG'] = 'L'
    aa_nucl_map['TTR'] = 'L'
    aa_nucl_map['YTR'] = 'L'

    aa_nucl_map['ATG'] = 'M'

    aa_nucl_map['AAC'] = 'N'
    aa_nucl_map['AAT'] = 'N'
    aa_nucl_map['AAY'] = 'N'

    four_fold_deg('CC', 'P')


    aa_nucl_map['CAA'] = 'Q'
    aa_nucl_map['CAG'] = 'Q'
    aa_nucl_map['CAR'] = 'Q'


    four_fold_deg('CG', 'R')
    aa_nucl_map['AGA'] = 'R'
    aa_nucl_map['AGG'] = 'R'
    aa_nucl_map['AGR'] = 'R'
    aa_nucl_map['MGR'] = 'R'



    four_fold_deg('TC', 'S')
    aa_nucl_map['AGC'] = 'S'
    aa_nucl_map['AGT'] = 'S'
    aa_nucl_map['WGY'] = 'S'

    four_fold_deg('AC', 'T')

    four_fold_deg('GT', 'V')

    aa_nucl_map['TGG'] = 'W'

    aa_nucl_map['TAC'] = 'Y'
    aa_nucl_map['TAT'] = 'Y'
    aa_nucl_map['TAY'] = 'Y'

    aa_nucl_map['---'] = '-'
    aa_nucl_map['...'] = '.'

    return aa_nucl_map

aa_nucl_map = create_aa_nucl_map()
from numpy import *
import sys
aa_nucl_map_ = zeros([256,256,256], dtype =int8)
aa_nucl_map_.flat = ord('X')
for k,v in aa_nucl_map.iteritems():
    aa_nucl_map_[ord(k[0]),ord(k[1]),ord(k[2])] = ord(v)


#    return "".join([ code.get(cdna[i:i+3], "?")
#                    for i in xrange(0,len(cdna),3) ])


def translate (cdna):
    return "".join([aa_nucl_map.get(cdna[x:x+3], 'X')
                        for x in xrange(0, len(cdna), 3) ])

from numpy import *
def fast_translate (cdna):
    cdna_len = len(cdna) - 2
    peptide = zeros((cdna_len+2) // 3, dtype = int8)
    for i, x in enumerate(xrange(0, cdna_len, 3)):
        peptide[i] = aa_nucl_map_[ord(cdna[x]), ord(cdna[x+1]), ord(cdna[x+2])]
    return peptide.tostring()




#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
from time_function import print_timing
class Test_translate(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_function(self):
        """
            test
        """
        def unit_test_help (cdna):
            cdna_len = len(cdna)
            #print "\t", cdna_len, cdna_len // 3, cdna_len % 3, len(translate(cdna)), translate(cdna)
            #print "\t", cdna_len, cdna_len // 3, cdna_len % 3, len(fast_translate(cdna)), fast_translate(cdna)
            return fast_translate(cdna)

        cdna = ("CTGCAATTGTACCACTGGGTGGAAAGGATCATTCTGCACAGAGACAGTTTCTACCTGTGATCCTGAACATGACCC"
                "TCCACACCACTGTAGCAGAGGAGCAACCTGCATTTCATTACCTCATGGATACACCTGTTTCTGTCCTCTAGGA"
                "ACCACTGGAATCTACTGTGAACAAGR")
        cdna = cdna.upper() + "  "
        self.assertEqual(unit_test_help(cdna),
                        "LQLYHWVERIILHRDSFYL*S*T*PSTPL*QRSNLHFITSWIHLFLSSRNHWNLL*TR")
        self.assertEqual(unit_test_help(cdna[1:]),
                        "CNCTTGWKGSFCTETVSTCDPEHDPPHHCSRGATCISLPHGYTCFCPLGTTGIYCEQX")
        self.assertEqual(unit_test_help(cdna[2:]),
                        "AIVPLGGKDHSAQRQFLPVILNMTLHTTVAEEQPAFHYLMDTPVSVL*EPLESTVNKX")


    def test_big(self):
        """
            test
        """
        import random

        @print_timing
        def fast_function (cdna):
            cdna_len = len(cdna)
            for a in range(5):
                res =  fast_translate(cdna)
            print "\t", cdna_len, cdna_len // 3, cdna_len % 3, len(res)
            return res

        @print_timing
        def slow_function (cdna):
            cdna_len = len(cdna)
            for a in range(5):
                res =  translate(cdna)
            print "\t", cdna_len, cdna_len // 3, cdna_len % 3, len(res)
            return res

        cdna = list()
        print >>sys.stderr,  "create_array"
        for a in xrange(1000):
            cdna.append(random.choice('ACGT'))
        cdna = "".join(cdna).upper() + "  "
        print >>sys.stderr,  "translate fast"
        fast_function(cdna)
        print >>sys.stderr,  "translate slow"
        slow_function(cdna)

#
#   debug code not run if called as a module
#
if __name__ == '__main__':
    import sys
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()




